﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using AzureMRCore.DataModel;
using AzureMRCore.IterativeDataCache;
using AzureMRCore.Utils;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.ServiceRuntime;
using Microsoft.WindowsAzure.StorageClient;

namespace AzureMRCore.MapRed
{
    public class MapWorker
    {
        private DataCache _dataCache;
        private CloudQueue _mapQueue;
        private MRConf _mrConf;
        private CloudStorageAccount _storageAccount;
        private string _tableAddress;
        private StorageCredentials _storageCredentials;
        private CloudTableClient _tableClient;
        private Guid _workerId;
        private string _roleId;
        private Dictionary<string, List<MapTaskHistory>> _mapTaskHistories;

        public MapWorker(CloudStorageAccount storageAccount, CloudQueue queue, MRConf mrConf, DataCache dataCache,
                         Dictionary<string, List<MapTaskHistory>> mapTaskHistories)
        {
            _mrConf = mrConf;
            _storageAccount = storageAccount;
            _tableAddress = _storageAccount.TableEndpoint.ToString();
            _mapQueue = queue;
            _dataCache = dataCache;
            _storageCredentials = _storageAccount.Credentials;
            _mapTaskHistories = mapTaskHistories;

            _tableClient = _storageAccount.CreateCloudTableClient();

            _workerId = Guid.NewGuid();
            _roleId = RoleEnvironment.CurrentRoleInstance.Id;
        }

        public void RunMapWorker()
        {
            var jobBoardDataContext = new JobBoardDataContext(_tableAddress, _storageCredentials,
                                                              DataModelUtils.GetJobBoardTableName());
            //jobBoardDataContext.MergeOption = MergeOption.NoTracking;
            _tableClient.CreateTableIfNotExist(jobBoardDataContext.Name);

            int countLimit = 4;
            string configString = RoleEnvironment.GetConfigurationSettingValue("CountLimit");
            if (configString != null & configString != "")
            {
                countLimit = Int32.Parse(configString);
            }

            // when job bulleting is not empty, give a hueristic priority to cached tasks with a time buffer before going to queue
            int cachePriority = 0;
            int count = 0;
            while (true)
            {
                try
                {
                    Stopwatch mapTimer = Stopwatch.StartNew();

                    bool taskFromCache = false;

                    /* First check the task table*/
                    var jobEnum = jobBoardDataContext.Tasks;
                    cachePriority++;

                    foreach (var job in jobEnum)
                    {
                        JobBoardDataModel iterativeJob = job;

                        if (iterativeJob.IsPrimitiveBCast &&
                            !(_mrConf.BCastReceiver.ContainsBCastData(iterativeJob.BCastDataURI)))
                        {
                            continue;
                        }
                        //todo: do the same for disk cache
                        List<MapTaskHistory> mapTaskHistories = GetMapTaskHistories(iterativeJob);
                        if (mapTaskHistories == null)
                        {
                            continue;
                        }

                        for (int i = 0; i < mapTaskHistories.Count; ++i)
                        {
                            MapTaskHistory mapTaskHistory = mapTaskHistories[i];
                            CachedItem item = null;
                            if (_dataCache.InMemCache.TryGetValue(mapTaskHistory.BlobURL, out item))
                            {
                                if (mapTaskHistory.AddChild(mapTaskHistory.JobID, mapTaskHistory.AppName,
                                                            iterativeJob.RowKey))
                                {
                                    var strings = mapTaskHistory.MapID.Split('_');
                                    string mapIndex = strings[0];
                                    string mapTableName = DataModelUtils.GetMapTableName(iterativeJob.JobID);
                                    var mapDContext = new MapMetaDataContext(_tableAddress, _storageCredentials,
                                                                             mapTableName);
                                    var mapID = mapIndex + "_" + iterativeJob.AppName;
                                    var mapTaskEntity = new MapDataModel(mapID,
                                                                         iterativeJob.RowKey, iterativeJob.AppName,
                                                                         mapTaskHistory.BlobURL)
                                                            {
                                                                Attempt1Status = MapDataModel.SUBMITTED,
                                                                Attempt2Status = MapDataModel.PENDING,
                                                                Attempt3Status = MapDataModel.PENDING,
                                                                TaskStatus = MapDataModel.SUBMITTED,
                                                                BroadcastData = iterativeJob.BCastDataURI,
                                                                NumMapTasks = mapTaskHistory.NumTotalMapTasks
                                                            };

                                    mapDContext.AddObject(mapDContext.Name, mapTaskEntity);
                                    try
                                    {
                                        mapDContext.SaveChanges();
                                    }
                                    catch (Exception e)
                                    {
                                        //TODO use blob lease to acquire a lock for this purpose
                                        // if creating map task in table errors, then ignore this task and move on. It may be processed by another worker
                                        mapTaskHistory.RemoveChild(mapTaskHistory.JobID, mapTaskHistory.AppName,
                                                                   iterativeJob.RowKey);
                                        continue;
                                    }
                                    var newMapHistory = new MapTaskHistory
                                                            {
                                                                MapID = mapID,
                                                                AppName = iterativeJob.AppName,
                                                                JobID = iterativeJob.JobID,
                                                                Iteration = iterativeJob.RowKey,
                                                                BlobURL = mapTaskHistory.BlobURL,
                                                                NumTotalMapTasks = mapTaskHistory.NumTotalMapTasks
                                                            };

                                    AddMapTaskHistory(newMapHistory);

                                    taskFromCache = true;
                                    count = 0;
                                    bool success = ExecuteTaskFromCache(iterativeJob, item, mapTaskHistory.BlobURL,
                                                                        mapTaskEntity, mapDContext, mapTimer);

                                    mapTimer.Restart();
                                }
                            }
                        }

                        /*MapMetaDataContext mapDataContext = new MapMetaDataContext(_tableAddress,
                                                                                   _storageCredentials,
                                                                                   DataModelUtils.GetMapTableName(
                                                                                       iterativeJob.JobID));

                        //we are not using this context for any updates/additions
                        mapDataContext.MergeOption = MergeOption.NoTracking;

                        IQueryable<MapDataModel> mapTasks = mapDataContext.MapTasks.Where(p =>
                                                                                          (p.PartitionKey ==
                                                                                           iterativeJob.RowKey.ToString() &
                                                                                           p.TaskStatus !=
                                                                                           MapDataModel.EXECUTING &
                                                                                           p.TaskStatus !=
                                                                                           MapDataModel.FINISHED));

                        if (!mapTasks.GetEnumerator().MoveNext()) continue;

                        //                        Trace.WriteLine("DataCache size:" + _dataCache.InMemCache.Count + " WorkerID:" +
                        //                                        _workerId);

                        //TODO: Can use the inDir container to short list blobs
                        IEnumerable<CachedItem> cachedItems = _dataCache.InMemCache.Where(entry =>
                                                                                          (entry.Value.JobID ==
                                                                                           iterativeJob.JobID
                                                                                          )).Select(entry => entry.Value);
                        //& entry.Value.Iteration < iterativeJob.RowKey


                        // using a list to avoid "collection modification exceptions" while doing the enumeration
                        var cachedItemsList = cachedItems.ToList();
                        foreach (CachedItem item in cachedItemsList)
                        {
                            string itemBlobURL = item.BlobURL;
                            IQueryable<MapDataModel> mapTaskEnum =
                                mapTasks.Where(p => (p.BlobURL == itemBlobURL));
                            MapDataModel mapTask = mapTaskEnum.FirstOrDefault();
                            if (mapTask != null)
                            {
                                taskFromCache = true;
                                bool success = ExecuteTaskFromCache(iterativeJob, item, itemBlobURL, mapTask, mapTimer);
                                mapTimer.Restart();
                                cachePriority = 0;
                            }
                        }*/
                    }

                    // Give priority to cached tasks
                    if ((!taskFromCache) & (count < countLimit)) // & (cachePriority == 0 || cachePriority > 4))
                    {
                        cachePriority = 0;
                        //Trace.WriteLine("Going to queue.e WorkerID:" +_workerId+"  "+_roleId);
                        //get the input message
                        CloudQueueMessage msg =
                            _mapQueue.GetMessage(new TimeSpan(_mrConf.MapTimeOut/3600, (_mrConf.MapTimeOut%3600)/60,
                                                              (_mrConf.MapTimeOut%60)));
                        if (msg != null)
                        {
                            Dictionary<string, object> mapParams = JSONUtils.parseJSON(msg.AsString);
                            var mapTaskContext = new MapTaskContext(mapParams, _storageAccount);

//                            Trace.WriteLine("Start Processing Map. MapTask:" + mapTaskContext.MapTask.PartitionKey + " " +
//                                            mapTaskContext.MapTask.RowKey + " WorkerID:" + _workerId + "  " + _roleId);

                            if (mapTaskContext.MapTask.TaskStatus == MapDataModel.FINISHED)
                            {
                                _mapQueue.DeleteMessage(msg);
                            }
                                // TODO handle executing messages
                            else if (mapTaskContext.MapTask.TaskStatus != MapDataModel.EXECUTING)
                            {
                                //better to throw exception early, if somebody else got this task
                                mapTaskContext.SetStateToExecute();

                                MapTaskHistory newMapHistory = new MapTaskHistory
                                                                   {
                                                                       MapID = mapTaskContext.MapID,
                                                                       AppName = mapTaskContext.AppName,
                                                                       JobID = mapTaskContext.JobID,
                                                                       Iteration = mapTaskContext.Iteration,
                                                                       BlobURL = mapTaskContext.MapTask.BlobURL,
                                                                       NumTotalMapTasks =
                                                                           mapTaskContext.MapTask.NumMapTasks
                                                                   };

                                AddMapTaskHistory(newMapHistory);

                                var mrDriver = _mrConf.MRDrivers[mapTaskContext.AppName];
                                IMapper mapper = mrDriver.GetMapper();
                                mapper.Init(_storageAccount, _mrConf, mrDriver, _dataCache, mrDriver.ConfigOut);
                                bool success = mapper.ExecuteMapTask(mapTaskContext);

                                if (success)
                                {
                                    try
                                    {
                                        _mapQueue.DeleteMessage(msg); //only if succesfull
                                    }
                                    catch
                                    {
                                        // Ignore the exception when message cannot be deleted
                                    }
                                }
                                mapTaskContext.MapTask.WorkerRole = _roleId;
                                mapTaskContext.RecordTimeAndSave(mapTimer);
                                mapTimer.Restart();
                                count++;
                            }
                        }
                        else
                        {
                            Thread.Sleep(_mrConf.MapQPollingInterval);
                            continue;
                        }
                    }
                    else if (!taskFromCache)
                    {
                        Thread.Sleep(_mrConf.MapQPollingInterval);
                        continue;
                    }
                }
                catch (Exception e)
                {
                    Trace.TraceError(e.ToString());
                    var errorContext = new ErrorDataContext(_tableAddress,
                                                            _storageCredentials, "AzureMRErrors");
                    _tableClient.CreateTableIfNotExist(errorContext.Name);
                    var error = new ErrorDataModel(e, "maptask", e.Message);
                    errorContext.AddObject(errorContext.Name, error);
                    errorContext.SaveChangesWithRetries();
                }
            }
        }

        private bool ExecuteTaskFromCache(JobBoardDataModel iterativeJob, CachedItem cachedDataItem, string itemBlobURL,
                                          MapDataModel mapTask, MapMetaDataContext mapDataContext, Stopwatch mapTimer)
        {
//            Trace.WriteLine("Start Processing Map. MapTask:" + mapTask.PartitionKey + " " +
//                            mapTask.RowKey + " WorkerID:" +
//                            _workerId + "  " + _roleId);
            //put these information in the jobboard too, so that we don't have to query the job too


            var mapTaskContext = new MapTaskContext(_storageAccount)
                                     {
                                         MapDataContext = mapDataContext,
                                         MapTask = mapTask,
                                         IsBlobType = true,
                                         AppName = mapTask.AppName,
                                         BlobUrl = itemBlobURL,
                                         Iteration = iterativeJob.RowKey,
                                         JobID = iterativeJob.JobID,
                                         MapID = mapTask.RowKey,
                                         NumReduceTasks = iterativeJob.NumReduceTasks,
                                         OutDir = iterativeJob.OutDir,
                                         ProgramParams = iterativeJob.ProgramParams,
                                         IsPrimitiveBCast = iterativeJob.IsPrimitiveBCast
                                     }; //TODO tasktype
            mapTaskContext.SetStateToExecute();

            var mrDriver = _mrConf.MRDrivers[mapTaskContext.AppName];
            IMapper mapper = mrDriver.GetMapper();
            mapper.Init(_storageAccount, _mrConf, mrDriver, _dataCache, mrDriver.ConfigOut);
            bool success = mapper.ExecuteMapTask(cachedDataItem.CachedInputFormat.GetCachedCopy(),
                                                 mapTaskContext);

            mapTaskContext.MapTask.ProcessedFromCache = true;
            mapTaskContext.MapTask.WorkerRole = _roleId;
            //Update the status of map task and in data cache  blah..blah...
            mapTaskContext.RecordTimeAndSave(mapTimer);
            return success;
        }

        private List<MapTaskHistory> GetMapTaskHistories(JobBoardDataModel iterativeJob)
        {
            List<MapTaskHistory> mapList;
            if (
                _mapTaskHistories.TryGetValue(
                    iterativeJob.JobID + "_" + iterativeJob.BaseAppName + "_" + iterativeJob.BaseIteration, out mapList))
            {
                return mapList;
            }
            return null;
        }

        private void AddMapTaskHistory(MapTaskHistory newMapHistory)
        {
            List<MapTaskHistory> mapList;
            string key = newMapHistory.JobID + "_" + newMapHistory.AppName + "_" + newMapHistory.Iteration;
            lock (_mapTaskHistories)
            {
                if (_mapTaskHistories.TryGetValue(key, out mapList))
                {
                    mapList.Add(newMapHistory);
                }
                else
                {
                    mapList = new List<MapTaskHistory> {newMapHistory};
                    _mapTaskHistories.Add(key, mapList);
                }
            }
        }
    }
}